#!/usr/bin/env python3

#  Imports.
import tmm.web.base.encoding as _encoding
import tmm.web.mq.bucket as _mq_bucket
import tmm.web.mq.subscription as _mq_subscription
import json as _json

#  Third-party libraries.
import tornado.websocket as _ws

BUCKET_SIZE = None


# noinspection PyAbstractClass,PyClassHasNoInit
class MessageQueueHandler(_ws.WebSocketHandler):
    """Message handler."""

    def open(self):
        """Event that will be triggered when the connection is opened."""

        #  Create a new bucket.
        # noinspection PyAttributeOutsideInit
        self._bucket = _mq_bucket.MessageBucket(limit=BUCKET_SIZE)

        #  Subscribe.
        _mq_subscription.add_subscriber(self._bucket)

    def on_message(self, message):
        """Event that will be triggered when receiving a message.

        :type message: bytes | str
        :param message: The message.
        """

        #  We don't deal with binary messages.
        if not isinstance(message, str):
            return

        #  Standardize.
        message = message.strip().lower()

        if message == "fetch":
            #  Fetch(read) messages from the bucket.
            raw_msgs = self._bucket.read_messages(clear=True)

            #  Pre-process.
            to_send = []
            for raw_msg in raw_msgs:
                try:
                    to_send.append(_json.loads(
                        raw_msg.decode(encoding=_encoding.TMM_WEB_ENCODING),
                        encoding=_encoding.TMM_WEB_ENCODING,
                    ))
                except ValueError:
                    pass

            #  Serialize and send.
            self.write_message(_json.dumps(to_send, indent=None))
        elif message == "shutdown":
            self.close()
        else:
            pass

    def on_close(self):
        """Event that will be triggered when the connection is closed."""

        #  Un-subscribe.
        _mq_subscription.remove_subscriber(self._bucket)
